#include "dpc.h"

Image* create_image(int width, int height) {
    Image* img = (Image*) malloc(sizeof(Image));
    img->data = (uint16_t*) calloc(width*height, sizeof(uint16_t));
    img->width = width;
    img->height = height;
    printf("create_image start with width=%d, height=%d\n", width, height);
    show_image(img, width, height);
    printf("create_image end\n");
    return img;
}

void destroy_image(Image* img) {
    free(img->data);
    free(img);
}

void show_image(Image* img, int width, int height) {
    for(int i=0; i<width*height; i++){
        if(i%width==0 && i!=0) printf("\n");
        printf("%3d ", img->data[i]);
    }
    printf("\n");
}
Image* pad_image(Image* img, int pad_size) {
    // [1 2 3]
    // [4 5 6]
    // [7 8 9]
    // np.pad(img, ((2, 2), (2, 2)), 'reflect')
    // [9 8 7 8 9 8 7]
    // [6 5 4 5 6 5 4]
    // [3 2 1 2 3 2 1]
    // [6 5 4 5 6 5 4]
    // [9 8 7 8 9 8 7]
    // [6 5 4 5 6 5 4]
    // [3 2 1 2 3 2 1]
    int new_width = img->width + pad_size*2;
    int new_height = img->height + pad_size*2;
    Image* padded_image = create_image(new_width, new_height);

    for(int y=pad_size; y<new_height-pad_size; y++) {
        for(int x=pad_size-1; x>=0; x--) {
            padded_image->data[y*new_width+x] = img->data[(y-pad_size)*img->width + (pad_size-x)];
        }
        for(int x=pad_size; x<pad_size+img->width; x++) {
            padded_image->data[y*new_width+x] = img->data[(y-pad_size)*img->width + x-pad_size];
        }
        for(int x=pad_size+img->width; x<padded_image->width; x++) {
            padded_image->data[y*new_width+x] = img->data[(y-pad_size+1)*img->width -2 -(x-(pad_size+img->width))];
        }
    }

    for(int y=0; y<pad_size; y++) {
        for(int x=0; x<new_width; x++) {
            padded_image->data[y*new_width+x] = padded_image->data[(pad_size+pad_size-y)*new_width+x];
        }
    }

    for(int y=new_height-pad_size; y<new_height; y++) {
        for(int x=0; x<padded_image->width; x++) {
            padded_image->data[y*new_width+x] = padded_image->data[(new_height-pad_size-(2+y-(new_height-pad_size)))*new_width+x];
        }
    }
    return padded_image;
}

void clip_image(Image *img, uint16_t minVal, uint16_t maxVal) {
    for (int i = 0; i < img->width * img->height; i++) {
        img->data[i] = MIN(MAX(img->data[i], minVal), maxVal);
    }
}

void dead_pixel_correction(Image* img, int thres, Mode_e mode, uint16_t clip) {
    Image *paddedImg = pad_image(img, 2);
    Image *dpcImg = create_image(img->width, img->height);

    for (int y = 2; y < paddedImg->height - 2; y++) {
        for (int x = 2; x < paddedImg->width - 2; x++) {
            int p0 = paddedImg->data[y * paddedImg->width + x];
            int p1 = paddedImg->data[(y - 2) * paddedImg->width + (x - 2)];
            int p2 = paddedImg->data[(y - 2) * paddedImg->width + x];
            int p3 = paddedImg->data[(y - 2) * paddedImg->width + (x + 2)];
            int p4 = paddedImg->data[y * paddedImg->width + (x - 2)];
            int p5 = paddedImg->data[y * paddedImg->width + (x + 2)];
            int p6 = paddedImg->data[(y + 2) * paddedImg->width + (x - 2)];
            int p7 = paddedImg->data[(y + 2) * paddedImg->width + x];
            int p8 = paddedImg->data[(y + 2) * paddedImg->width + (x + 2)];
            if (fabs(p1 - p0) > thres && fabs(p2 - p0) > thres && fabs(p3 - p0) > thres &&
                fabs(p4 - p0) > thres && fabs(p5 - p0) > thres && fabs(p6 - p0) > thres &&
                fabs(p7 - p0) > thres && fabs(p8 - p0) > thres) {
                if (mode == 0) {
                    p0 = (p2 + p4 + p5 + p7) / 4;
                } else if (mode == 1) {
                    int dv = fabs(2 * p0 - p2 - p7);
                    int dh = fabs(2 * p0 - p4 - p5);
                    int ddl = fabs(2 * p0 - p1 - p8);
                    int ddr = fabs(2 * p0 - p3 - p6);
                    int minDiff = MIN(MIN(MIN(dv, dh), ddl), ddr);
                    if (minDiff == dv) {
                        p0 = (p2 + p7 + 1) / 2;
                    } else if (minDiff == dh) {
                        p0 = (p4 + p5 + 1) / 2;
                    } else if (minDiff == ddl) {
                        p0 = (p1 + p8 + 1) / 2;
                    } else {
                        p0 = (p3 + p6 + 1) / 2;
                    }
                }
            }
            dpcImg->data[(y - 2) * img->width + (x - 2)] = (uint16_t)p0;
        }
    }
    memcpy(img->data, dpcImg->data, img->width * img->height * sizeof(uint16_t));
    destroy_image(paddedImg);
    destroy_image(dpcImg);
    clip_image(img, 0, 255); // Clip the image values
}

